module AnnotatingTypeChecker where

import AbsJavalette
import PrintJavalette
import ErrM
import Control.Monad(foldM)
import qualified Data.Map as Map
import Data.Map (Map, empty, insert, fromList)

type Env = (Sigs, [Context])
type Sigs = Map Ident ([Type], Type)
type Context = Map Ident Type
type TAEE = (Expr, Env)
type TASE = (Stmt, Env)

getExp :: TAEE -> Expr
getExp (expr, _) = expr

getEnv :: TAEE -> Env
getEnv (_, environment) = environment

--Helper to get the type from a type annotated expression
getType :: Expr -> Err Type
getType x = case x of
	ETyped t x -> return t
	_ -> fail "Not of type ETyped"

emptyEnv :: Env
emptyEnv = (sigs, context) where
	sigs = fromList [(Ident "printInt", ([Int], Void)), (Ident "printString", ([String], Void)), (Ident "printDouble", ([Doub], Void)), (Ident "readInt", ([], Int)), (Ident "readDouble", ([], Doub))]
	context = []

lookupVar :: Env -> Ident -> Err Type
lookupVar (s, []) identifier = fail ("No such variable " ++ show identifier)
lookupVar (s, curr:rest) identifier = case Map.lookup identifier curr of
	Nothing -> lookupVar (s, rest) identifier
	Just t -> return t

--Takes an environment and an identifier of a function and returns 
--the functions signature if exists or an error
lookupFun :: Env -> Ident -> Err ([Type], Type)
lookupFun (signatures, _) identifier = case Map.lookup identifier signatures of
	Nothing -> fail ("function not defined " ++ show identifier)
	Just (args, retType) -> return (args, retType)

--top is the top of the context stack, that is, the current scope
--TODO error if variable is redeclared, if exists in scope then error
updateVar :: Env -> Ident -> Type -> Err Env
updateVar (sigs, top:rest) identifier typ = case Map.lookup identifier top of
	Nothing -> return (sigs, (insert identifier typ top):rest)
	Just t -> fail "Redeclaration of variable in same scope"

collectTypes :: [Arg] -> [Type]
collectTypes [] = []
collectTypes ((Arg t _):rest) = t:(collectTypes rest)

updateFun :: Env -> TopDef -> Err Env
updateFun (sigs, context) (FnDef retType identifier args _) = 
	return (insert identifier (collectTypes args, retType) sigs, context)

--When entering a new block, a new empty context is pushed on the stack
newBlock:: Env -> Env
newBlock (sigs, context) = (sigs, empty:context)

checkExp2 :: Env -> Type -> Expr -> Err Expr
checkExp2 env typ expr = case inferExp2 env expr of
	Ok expr' -> case getType expr' of
		Ok t -> if(t == typ)
			then 
				return expr'
			else
				fail $ "Type of " ++ printTree expr ++ ", expected '" ++ printTree typ ++ "'" ++ " but found '" ++ printTree t ++ "'"
		Bad s -> fail s
	Bad s -> fail s

checkStmt2 :: Type -> Env -> Stmt -> Err (Env, Stmt)
checkStmt2 typ env stmt = case stmt of
	Empty -> return (env, stmt)
	BStmt (Block statements) -> do
		case blockChecker typ ((newBlock env), []) statements of
			Ok (_, stmts') -> return (env, (BStmt (Block stmts'))) 
			Bad s -> fail s
		where
			blockChecker :: Type -> (Env, [Stmt]) -> [Stmt] -> Err (Env, [Stmt])
			blockChecker t (env', stmts) [] = return (env', stmts)
			blockChecker t (env', stmts) (st:rest) = case checkStmt2 t env' st of
				Ok (e, st') -> blockChecker t (e, stmts++[st']) rest
				Bad s -> fail s
	Decl t items -> do

		case foldM addItems (env, []) items of
			Ok (env', items') -> return (env', Decl t items')
			Bad s -> fail s
		where
			addItems :: (Env, [Item]) -> Item -> Err (Env, [Item])
			addItems (en, is) i@(NoInit ident) = case updateVar en ident t of
				Ok en' -> return (en', is++[i])
				Bad s -> fail s
			addItems (en, is) i@(Init ident expr) = case checkExp2 en t expr of
				Ok expr' -> case (updateVar en ident t) of
					Ok en'' -> return (en'', is++[Init ident expr'])
					Bad s -> fail s
				Bad s -> fail s
	Ass identifier expr -> case lookupVar env identifier of
		Ok t -> case checkExp2 env t expr of
			Ok x -> return (env, (Ass identifier x))
			Bad s -> fail s
		Bad s -> fail s
	Incr identifier -> case lookupVar env identifier of
		Ok Int -> return (env, stmt)
		Ok t -> fail "No incrementation of non ints"
		Bad s -> fail s
	Decr identifier -> case lookupVar env identifier of
		Ok Int -> return (env, stmt)
		Ok t -> fail "No decrementation of non ints"
		Bad s -> fail s
	Ret expr -> case checkExp2 env typ expr of
		Ok x -> return (env, (Ret x))
		Bad s -> fail s
	VRet -> if typ == Void
		then return (env, stmt)
		else fail "Non void return"
	Cond expr stmt2 -> case checkExp2 env Bool expr of
		Ok x -> case checkStmt2 typ (newBlock env) stmt2 of
			Ok (_, stmt3) -> return (env, (Cond x stmt3))
			Bad s -> fail s
		Bad s -> fail s
	CondElse expr stmt2 stmt3 -> case checkExp2 env Bool expr of
		Ok x -> case checkStmt2 typ (newBlock env) stmt2 of
			Ok (_, stmt2') -> case checkStmt2 typ (newBlock env) stmt3 of
				Ok (_, stmt3') -> return (env, (CondElse x stmt2' stmt3'))
				Bad s -> fail s
			Bad s -> fail s
		Bad s -> fail s 
	While expr stmt2 -> case checkExp2 env Bool expr of
		Ok x -> case checkStmt2 typ (newBlock env) stmt2 of
			Ok (_, stmt2') -> return (env, (While x stmt2'))
			Bad s -> fail s
		Bad s -> fail (s ++ "Non boolean expression in while")
	SExp expr -> case inferExp2 env expr of
		Ok x -> return (env, (SExp x))
		Bad s -> fail s

checkReturnProgram :: Program -> Err ()
checkReturnProgram (Program topdefs) = case (foldr (&&) True (map checkReturnTopDef topdefs)) of
	True -> return ()
	False -> fail "Non reacheable return"

checkReturnTopDef :: TopDef -> Bool
checkReturnTopDef (FnDef Void _ _ _) = True
checkReturnTopDef (FnDef _ _ _ block) = checkReturnStmt (BStmt block)

checkReturnStmt :: Stmt -> Bool
checkReturnStmt (BStmt (Block stmts)) = foldr (||) False (map checkReturnStmt stmts)
checkReturnStmt (Cond ELitTrue stmt) = checkReturnStmt stmt
checkReturnStmt (CondElse ELitTrue stmt _) = checkReturnStmt stmt
checkReturnStmt (CondElse ELitFalse _ stmt) = checkReturnStmt stmt
checkReturnStmt (CondElse _ stmt1 stmt2) = (checkReturnStmt stmt1) && (checkReturnStmt stmt2)
checkReturnStmt (While ELitTrue stmt) = checkReturnStmt stmt
checkReturnStmt (Ret expr) = True
checkReturnStmt VRet = True
checkReturnStmt _ = False

updateEnvWithArg :: Env -> Arg -> Err Env
updateEnvWithArg environment (Arg typ identifier) = case updateVar environment identifier typ of
	Ok e -> return e
	Bad s -> fail s

block2Stmts :: Block -> [Stmt]
block2Stmts (Block stmts) = stmts

--Any return statements should be of the same type as the function
--Extend the environment with arguments and types
checkDef2 :: Env -> TopDef -> Err (Env, TopDef)
checkDef2 env fun@(FnDef typ identifier args block) = case foldM updateEnvWithArg (newBlock env) args of
	Ok env' -> do
		--checkReturn fun
		case checkStmt2 typ env' (BStmt block) of
			Ok (_, (BStmt block')) -> return (env, (FnDef typ identifier args block'))
			Bad s -> fail s
	Bad s -> fail s

checkProgram2 :: Program -> Err Program
checkProgram2 p@(Program fdefs) = case foldM updateFun (emptyEnv) fdefs of
	Ok e -> case checkReturnProgram p of
		Ok _ -> case functionChecker e [] fdefs of
			Ok fdefs' -> return (Program fdefs')
			Bad s -> fail s
		Bad s -> fail s
	Bad s -> fail s

functionChecker :: Env -> [TopDef] -> [TopDef] -> Err [TopDef]
functionChecker env defs [] = return defs
functionChecker env defs (d:rest) = case checkDef2 env d of
	Ok (_, d') -> functionChecker env (defs++[d']) rest
	Bad s -> fail s

--Type annotating version of inferExp
inferExp2 :: Env -> Expr -> Err Expr
inferExp2 env x = case x of
	ELitTrue -> return (ETyped Bool ELitTrue)
	ELitFalse -> return (ETyped Bool ELitFalse)
	ELitInt _ -> return (ETyped Int x)
	ELitDoub _ -> return (ETyped Doub x)
	EString _ -> return (ETyped String x)
	EVar identifier -> case lookupVar env identifier of
		Ok t -> return (ETyped t x)
		Bad s -> fail s
	Neg expression -> case inferExp2 env expression of --maybe we should check that the type of expression is of certain allowed types only, perhaps it is the case that not all types are negatable
		Ok e -> case getType e of
			Ok t -> return (ETyped t (Neg e))
			Bad s -> fail s
		Bad s -> fail s
	Not expression -> case inferExp2 env expression of
		Ok e -> case getType e of
			Ok t -> return (ETyped t (Not e))
			Bad s -> fail s
		Bad s -> fail s
	EMul leftE op rightE -> case inferExp2 env leftE of
		Ok le -> case getType le of --all theese case of Ok thingies can probably be removed if we acually use the monadic behaviour of Err, i guess its legal to just stack them on top of each other and check for errors in the end instead
			Ok t -> case checkExp2 env t rightE of
				Ok r -> return (ETyped t (EMul le op r))
				Bad s -> fail s
			Bad s -> fail s
		Bad s -> fail s
	EAdd leftE op rightE -> case inferExp2 env leftE of
		Ok le -> case getType le of
			Ok t -> case checkExp2 env t rightE of
				Ok r -> return (ETyped t (EAdd le op r))
				Bad s -> fail s
			Bad s -> fail s
		Bad s -> fail s
	ERel leftE op rightE -> case inferExp2 env leftE of
		Ok le -> case getType le of
			Ok t -> case checkExp2 env t rightE of
				Ok r -> return (ETyped Bool (ERel le op r))
				Bad s -> fail s
			Bad s -> fail s
		Bad s -> fail s
	EAnd leftE rightE -> case checkExp2 env Bool leftE of
		Ok le -> case checkExp2 env Bool rightE of
			Ok re -> return (ETyped Bool (EAnd le re))
			Bad s -> fail s
		Bad s -> fail s
	EOr leftE rightE -> case checkExp2 env Bool leftE of
		Ok le -> case checkExp2 env Bool rightE of
			Ok re -> return (ETyped Bool (EOr le re))
			Bad s -> fail s
		Bad s -> fail s
	EApp identifier expressions -> case lookupFun env identifier of
		Ok sig -> case checkFunctionArguments env [] (fst sig) expressions of
			Ok expressions' -> return (ETyped (snd sig) (EApp identifier expressions'))
			Bad s -> fail s
		Bad s -> fail s
	ETyped t expr -> return x

checkFunctionArguments :: Env -> [Expr] -> [Type] -> [Expr] -> Err [Expr]
checkFunctionArguments _ args' [] [] = return args'
checkFunctionArguments _ args' _ [] = fail "To few arguments in function call"
checkFunctionArguments _ args' [] _ = fail "To many arguments in function call"
checkFunctionArguments env args' (rt:rtRest) (ae:aeRest) = case checkExp2 env rt ae of
	Ok e -> checkFunctionArguments env (args'++[e]) rtRest aeRest
	Bad s -> fail (s++" in function call")
